#!/usr/bin/env ruby
## input: a fasta file from shotgun sequencing

# Genome size estimator for microbes
# Version 0.9.1 

# Author: Yufeng Shen, PhD
##        Human Genome Sequencing Center
#         Baylor College of Medicine
#         Houston, TX, 77030
#         ys135495@bcm.edu
##

## What's new:
#  0.9.1   use f(0) prediction as the estimate


## Todos:
#  1. deal with repeats and sequence errors





require 'getoptlong'

opts = GetoptLong.new(
    ["--fasta", "-f", GetoptLong::REQUIRED_ARGUMENT],
    ["--minSize", "-m", GetoptLong::OPTIONAL_ARGUMENT],
    ["--maxSize", "-x", GetoptLong::OPTIONAL_ARGUMENT],
	["--mvalue", "-v", GetoptLong::OPTIONAL_ARGUMENT],
	["--atlas","-a", GetoptLong::OPTIONAL_ARGUMENT],
    ["--help", "-h", GetoptLong::NO_ARGUMENT]
)

optHash = {}
opts.each do |opt, arg|
  optHash[opt] = arg
end

if optHash.key?("--help") or !optHash.key?("--fasta")
	$stderr.puts "\nUsage: ruby __.rb -f foo.fasta [-m min_read_size -x max_read_size -v mismatch -a atlas-overlapper_binary ]"
	$stderr.puts "\n    -f: input fasta file; ideally it should cover the genome 0.4x - 1.0x."
	$stderr.puts "        The fasta file should be trimmed.  For 454 FLX, trim 50bp off 3' end."
	$stderr.puts "\n OPTIONAL parameters: "
	$stderr.puts "    -m [150]: minimum read size; optional;  reads shorter than this will be discarded; default value: 150"
	$stderr.puts "    -x [320]: maximum read size; optional; default value: 320"
	$stderr.puts "    -v [5]: mismatch percentage allowed in reads overlap detection. (-m value for atlas-overlapper); default: 5 "
	$stderr.puts "    -a: the atlas-overlapper binary. Only required if your platform is not supported (neither Mac OS X Intel or Linux x86). "
  
	exit
end


# input fasta file
fa = optHash["--fasta"]

if optHash.key?("--mvalue")
	$mvalue = optHash["--mvalue"].to_i
else 
	$mvalue = 5
end

if optHash.key?("--minSize")
	$minReadSize = optHash["--minSize"].to_i
else
	$minReadSize = 150  # minimum read size
end

if optHash.key?("--maxSize")
	$maxReadSize = optHash["--maxSize"].to_i
else
	$maxReadSize = 320
end

$minoverlap = 24
$reads = {}

if optHash.key?("--atlas")
	$atlasOverlapper = optHash["--atlas"]
else
	os = `uname`.chomp
	binary = "atlas-overlapper-" + os
	
# assume the atlas-binner binary is under the same directory as this ruby engine
	$atlasOverlapper = File.dirname($0) + "/Atlas/"+binary  

	if !File.exist?($atlasOverlapper)
		$stderr.puts "The #{os} platform is not supported"
		exit
	end

end

if fa==nil or File.size(fa) < 10
  $stderr.puts "bad input file"
  exit
end


# build a overlap graph from the file
graph = fa + '.Intermediate.overlap_graph'

cmd = "#{$atlasOverlapper} -s #{fa} -q #{fa} -o #{graph} -m #{$mvalue} -G 1 -M #{$minReadSize} -R 2000 -H 2001 -b 0 -e 500000 -S 999 -B 7 -k 9 -I -O #{$minoverlap} -p 141999941 >/dev/null 2>&1"
system(cmd)


# parse the overlap graph
File.new(graph, 'r').each do |line|
  cols = line.split(/\t/)
  query = cols.shift
  if !$reads.key?(query)
    $reads[query] = {}
    $reads[query][:size] = 0
    $reads[query][:overlap] = {}
  end

  cols.each do |info|
    if info =~ /^(\S+)\s+\{(\d+)\,(\d+)\,(\S+)\,(\S+)\,([f|r])\,(\d+)/
      target,span,score = $1,$2.to_i, $3.to_i
      if target == query # 
        $reads[query][:size] = span
      else
        $reads[query][:overlap][target] = 1
        if !$reads.key?(target)
          $reads[target] = {}
          $reads[target][:size] = 0
          $reads[target][:overlap] = {}
        end
        $reads[target][:overlap][query] = 1
      end
    end
  end
end

## remove short or large reads                                               
avgSize = 0
$reads.each_key do |r|
  if $reads[r][:size] < $minReadSize or $reads[r][:size] > $maxReadSize
    $reads[r][:overlap].each_key do |t|
      $reads[t][:overlap].delete(r)
    end
    $reads[r].delete(:overlap)
    $reads.delete(r)
  else
    avgSize += $reads[r][:size]
  end
end
avgSize = avgSize / $reads.size.to_f


## count the overlap distribution                                                                        
oarray = (0..30).to_a
$overlap = {}
oarray.map {|i| $overlap[i]=0}

$reads.each_key do |r|
  o = $reads[r][:overlap].size
  
  if $overlap.key?(o)
    $overlap[o] += 1
  end
end


## estimate based on N(0),  N(1) / N(0), etc
ratio = $minoverlap / avgSize
numReads = $reads.size
cov = []
cov[0]  = Math.log($overlap[0] / $reads.size.to_f) / (-2.0 * (1 - ratio))

1.upto(10) do |i|
  cov << ($overlap[i] / $overlap[i-1].to_f) * (i / ((1-ratio)*2.0))
end
gsize = []


0.upto(5) do |i|
  gsize << (avgSize * $reads.size / cov[i]).to_i
#  puts "#{i}\t#{cov[i]}\t#{gsize[i]}"
end

puts "\nInput fasta file:  #{fa}"
puts "Number of reads used in estimation ------- #{numReads}"
puts "Average read size  ----------------------- #{avgSize.to_i} bp"
puts "Genome size estimated by "
puts "       f(0) : #{gsize[0]}  \t coverage: #{(cov[0]*1000).to_i/1000.0}"
puts "  f(1)/f(0) : #{gsize[1]}  \t coverage: #{(cov[1]*1000).to_i/1000.0}"
puts "  f(2)/f(1) : #{gsize[2]}  \t coverage: #{(cov[2]*1000).to_i/1000.0}"

# compute the stderr of the three estimates

avg = (gsize[0] + gsize[1] + gsize[2]) / 3.0

delta = 0

0.upto(2) do |i|
  delta += (avg - gsize[i])**2.0
end

stdvar = ((delta/2)**0.5).to_i

#puts "standard variation among predictions:  #{stdvar}"
puts "TABBED\t#{gsize[0]}\t#{gsize[1]}\t#{gsize[2]}\tstdvar\t#{stdvar}"
puts " "
puts "Estimated genome size (bp)  #{gsize[0]}"
puts " "


exit


# discussions:
# this is a pipeline for microbe genomes
## A question: is there any way to utilize the graph structure? better than island reads at low coverage ? deal with repeats? 
##
##
# Lander-Waterman model:
# read length : L
# read number: N
# genome size: G
# ---> the probability of being a start or end of a read for each genomic position is :
##  2 * N*L /G = 2C
##  let C'= C*(1-m/L), then
# f(K) = N*(2C')^K *e^(-2C')/K!  {1}
# 
# Based on equation {1}, we get :
# 1. number of island reads (defined as reads with no overlapping reads):
# f(0) = N*e^-2C'
##

# 2. linear model, model repeats and sequencing errors in a linear factor that affect f(0):
## N(0) = f(0) + r*f(0) + e
## r is a linear factor determined by repeats and errors.
## e is error adjustment

# 3. simple model:
# f(2)/f(1) = C'  --> assume the impact of repeats/errors are proportional on f(2) and f(1)
# f(1)/f(0) = 2C' --> for microbes, f(0) is usually larger than expected because of errors, thus this often leads to underestimates of C' and over-estimates of Genome size


## The big question is how to use the standard variation among different estimates to:
# 1. detect severe sampling or errors in the sample, AND
# 2. divide the sample into smaller pieces and bootstrap? 
